from channels.generic.websocket import WebsocketConsumer
import threading
import paramiko
import json
from io import StringIO
from system_config.models import Credential

class SSHConsumer(WebsocketConsumer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.ssh_ip = None
        self.ssh_port = None
        self.credential_id = None
        self.ssh_channel = None
        self.ssh_username = None
        self.ssh_password = None

    # 处理客户端与WebSocket建立连接
    def connect(self):
        self.ssh_ip = self.scope['url_route']['kwargs']['ssh_ip']
        self.ssh_port = self.scope['url_route']['kwargs']['ssh_port']
        try:
            self.credential_id = self.scope['url_route']['kwargs']['credential_id']
        except Exception as e:
            self.ssh_username = self.scope['url_route']['kwargs']['ssh_username']
            self.ssh_password = self.scope['url_route']['kwargs']['ssh_password']
        self.accept()  # 接受与客户端建立WebSocket连接
        self.ssh_stream() # 创建SSH连接和交互式Shell会话通道

    def ssh_stream(self):
        self.send(bytes_data='建立SSH连接中...\r\n'.encode('utf-8')) # 字符串转字节（必须）
        try:
            # 创建SSH连接
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            try:
                credential = Credential.objects.get(id=self.credential_id)
                username = credential.username
                if credential.auth_mode == 1:
                    password = credential.password
                    ssh.connect(self.ssh_ip, self.ssh_port, username, password=password, timeout=5)
                else:
                    cache = StringIO(credential.private_key)
                    private_key = paramiko.RSAKey.from_private_key(cache)
                    ssh.connect(self.ssh_ip, self.ssh_port, username, pkey=private_key, timeout=5)
            except Exception as e:
                ssh.connect(self.ssh_ip, self.ssh_port, self.ssh_username, self.ssh_password, timeout=5)
            # 创建SSH通道
            self.ssh_channel = ssh.invoke_shell(term='xterm') # 创建交互式Shell会话，并指定伪终端的终端类型为xterm
            self.ssh_channel.transport.set_keepalive(30)
            # 创建线程对象并启动（循环读取数据发送给前端）
            thread = threading.Thread(target=self._loop_read)
            thread.start()
        except Exception as e:
            if "timed out" in str(e):
                # 通过WebSocket向客户端发送消息
                self.send(bytes_data=f'建立SSH连接超时！\r\n'.encode('utf-8'))
                self.close()
            else:
                self.send(bytes_data=f'建立SSH连接失败！错误: {e}\r\n'.encode('utf-8'))
                self.close()
            return

    # 循环读取SSH通道返回的数据并发送给前端
    def _loop_read(self):
        while True:
            data = self.ssh_channel.recv(32 * 1024)
            if not data:
                break
            self.send(bytes_data=data)

    # 处理从客户端接收到的 WebSocket 消息
    def receive(self, text_data=None, bytes_data=None):
        data = text_data or bytes_data
        if data:
            data = json.loads(data)
            # 动态设置终端窗口大小
            resize = data.get('resize')
            try:
                if resize and len(resize) == 2:
                    self.ssh_channel.resize_pty(*resize)
                else:
                    self.ssh_channel.send(data['data'])  # Shell命令
            except Exception as e:
                print('接收前端终端窗口数据异常，通常是SSH通道建立异常')

    # 关闭WebSocket连接时调用
    def disconnect(self, close_code):
        try:
            self.ssh_channel.close()
        except Exception as e:
            print('关闭SSH通道异常，通常是SSH通道建立异常')

